import os
import random

import numpy as np
import xml.etree.ElementTree as ET

from PIL import Image, ImageEnhance
import matplotlib.pyplot as plt
import matplotlib.patches as patches

import paddle

object_names = ['Boerner', 'Leconte', 'Linnaeus', 'acuminatus', 'armandi', 'coleoptera', 'linnaeus'] # 物体名称
def get_object_gtcls():
    """
    功能:
        将物体名称映射成物体类别
    输入:
    输出:
        object_gtcls - 物体类别
    """
    object_gtcls = {} # 物体类别字典
    for key, value in enumerate(object_names):
        object_gtcls[value] = key # 将物体名称映射成物体类别
        
    return object_gtcls

def get_data_list(data_path):
    """
    功能:
        读取数据列表
    输入:
        data_path - 数据路径
    输出:
        data_list - 数据列表
    """
    file_list = os.listdir(os.path.join(data_path, 'annotations', 'xmls')) # 文件列表
    data_list = [] # 数据列表
    
    for file_id, file_name in enumerate(file_list):
        # 读取数据文件
        file_path = os.path.join(data_path, 'annotations', 'xmls', file_name) # 文件路径
        xml_tree = ET.parse(file_path) # 解析文件
        
        # 读取图像数据
        image_path = os.path.join(data_path, 'images', file_name.split('.')[0] + '.jpeg') # 图像路径
        image_w = float(xml_tree.find('size').find('width').text) # 图像宽度
        image_h = float(xml_tree.find('size').find('height').text) # 图像高度
        
        # 读取物体数据
        object_list = xml_tree.findall('object') # 物体列表
        object_gtbox = np.zeros((len(object_list), 4), dtype=np.float32) # 物体边框
        object_gtcls = np.zeros((len(object_list),  ), dtype=np.int32) # 物体类别
        
        for object_id, object_item in enumerate(object_list):
            # 读取物体类别
            object_name = object_item.find('name').text # 读取物体名称
            object_gtcls[object_id] = get_object_gtcls()[object_name] # 将物体名称映射成物体类别
            
            # 读取物体边框
            x_min = float(object_item.find('bndbox').find('xmin').text)
            y_min = float(object_item.find('bndbox').find('ymin').text)
            x_max = float(object_item.find('bndbox').find('xmax').text)
            y_max = float(object_item.find('bndbox').find('ymax').text)
            
            x_min = max(0.0, x_min)
            y_min = max(0.0, y_min)
            x_max = min(x_max, image_w - 1.0)
            y_max = min(y_max, image_h - 1.0)
            
            object_gtbox[object_id] = [
                (x_min + x_max) / 2.0, (y_min + y_max) / 2.0,
                (x_max - x_min + 1.0), (y_max - y_min + 1.0)] # 计算物体边框: xywh格式
        
        # 保存数据列表
        data = {'image_path': image_path, 'image_width': image_w, 'image_height': image_h,
                'object_gtbox': object_gtbox, 'object_gtcls': object_gtcls}
        
        if len(data) != 0:
            data_list.append(data)
            
    return data_list

##############################################################################################################

def get_object_item(object_gtbox, object_gtcls, image_size):
    """
    功能:
        读取物体项目
    输入:
        object_gtbox - 物体边框
        object_gtcls - 物体类别
        image_size   - 图像高宽
    输出:
        gtbox        - 边框列表
        gtcls        - 类别列表
    """
    # 添加物体数据列表
    max_num = 50 # 最大物体数量
    gtbox = np.zeros((max_num, 4)) # 边框列表
    gtcls = np.zeros((max_num,  )) # 类别列表
    
    for i in range(len(object_gtbox)):
        gtbox[i, :] = object_gtbox[i, :]
        gtcls[i] = object_gtcls[i]
        
        if i > max_num: # 是否超过最大物体数量
            break
    
    # 转换为相对真实图像比例的边框
    gtbox[:, 0] = gtbox[:, 0]/float(image_size[1]) # 计算边框x相对值
    gtbox[:, 1] = gtbox[:, 1]/float(image_size[0]) # 计算边框y相对值
    gtbox[:, 2] = gtbox[:, 2]/float(image_size[1]) # 计算边框w相对值
    gtbox[:, 3] = gtbox[:, 3]/float(image_size[0]) # 计算边框h相对值
    
    return gtbox, gtcls

def get_data_item(data):
    """
    功能: 
        读取数据项目
    输入: 
        data       - 数据
    输出:
        image      - 图像数据
        gtbox      - 边框列表
        gtcls      - 类别列表
        image_size - 图像高宽
    """
    # 读取数据项目
    image_path = data['image_path'] # 图像路径
    image_w = data['image_width'] # 图像宽度
    image_h = data['image_height'] # 图像高度
    object_gtbox = data['object_gtbox'] # 物体边框
    object_gtcls = data['object_gtcls'] # 物体类别
    
    # 打开图像文件
    image = Image.open(image_path) # 打开图像
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    image = np.array(image) # 转换为ndarray格式, 读取数据类型为HWC,uint8类型
    
    # 检查图像高宽
    assert image.shape[0] == int(image_h), \
        'image path: {}, image.shape[0]:{} != image_h: {}'.format(
        image_path, image.shape[0], image_h)
    assert image.shape[1] == int(image_w), \
        'image path: {}, image.shape[1]:{} != image_w: {}'.format(
        image_path, image.shape[1], image_w)
    
    # 读取物体项目
    image_size = (image_h, image_w)
    gtbox, gtcls = get_object_item(object_gtbox, object_gtcls, image_size)
    
    return image, gtbox, gtcls, image_size

##############################################################################################################

def random_distort_image(image):
    """
    功能: 
        随机变换图像
    输入: 
        image - 图像数据
    输出:
        image - 图像数据
    """
    # 随机变换饱和度
    def random_distort_saturation(image):
        if random.random() > 0.5:
            delta = random.uniform(-0.5, 0.5) + 1
            image = ImageEnhance.Color(image).enhance(delta)
        return image
    
    # 随机变换明亮度
    def random_distort_brightness(image):
        if random.random() > 0.5:
            delta = random.uniform(-0.5, 0.5) + 1
            image = ImageEnhance.Brightness(image).enhance(delta)
        return image
    
    # 随机变换对比度
    def random_distort_constract(image):
        if random.random() > 0.5:
            delta = random.uniform(-0.5, 0.5) + 1
            image = ImageEnhance.Contrast(image).enhance(delta)
        return image
    
    # 随机变换色调
    def random_distort_hue(image):
        if random.random() > 0.5:
            delta = random.uniform(-18, 18)
            image = np.array(image.convert('HSV'))
            image[:, :, 0] = image[:, :, 0] + delta # 变换色调
            image = Image.fromarray(image, mode='HSV').convert('RGB')
        return image
    
    # 随机变换图像
    distort = [random_distort_saturation, random_distort_brightness, random_distort_constract, random_distort_hue] # 变换方法列表
    np.random.shuffle(distort) # 打乱变换顺序
    
    image = Image.fromarray(image) # 转换为Image格式
    image = distort[0](image)
    image = distort[1](image)
    image = distort[2](image)
    image = distort[3](image)
    image = np.asarray(image) # 转换为ndarray格式
    
    return image

def random_expand_image(image, gtbox, keep_ratio=True):
    """
    功能: 
        随机填充图像
    输入: 
        image      - 图像数据
        gtbox      - 边框列表
        keep_ratio - 保持宽高比例
    输出:
        image      - 图像数据
        gtbox      - 边框列表
    """
    # 是否填充图像
    if random.random() > 0.5:
        return image, gtbox
    
    # 生成填充比例
    max_ratio = 4
    x_ratio = random.uniform(1, max_ratio) # 随机产生x填充比例
    
    if keep_ratio: # 是否保持宽高比例
        y_ratio = x_ratio # y填充比例等于x填充比例
    else:
        y_ratio = random.uniform(1, max_ratio) # 随机产生y填充比例
        
    # 计算填充宽高
    image_h, image_w, image_channel = image.shape # 获取图像高宽和通道数
    
    expand_w = int(image_w * x_ratio) # 计算填充宽度
    expand_h = int(image_h * y_ratio) # 计算填充高度
    
    x_offset = random.randint(0, expand_w - image_w) # 随机生成原图在填充图中x坐标位置
    y_offset = random.randint(0, expand_h - image_h) # 随机生成原图在填充图中y坐标位置
    
    # 生成填充图像
    expand_image = np.zeros((expand_h, expand_w, image_channel)) # float32类型
    
    mean_value = [0.485, 0.456, 0.406] # COCO数据集通道平均值
    for i in range(image_channel):
        expand_image[:, :, i] = mean_value[i] * 255.0 # 使用均值填充每个通道
    
    # 填充原始图像
    expand_image[y_offset : y_offset + image_h, x_offset : x_offset + image_w, :] = image # 填充原始图像
    image = expand_image.astype('uint8') # 转换为uint8类型
    
    # 计算相对边框
    gtbox[:, 0] = ((gtbox[:, 0] * image_w) + x_offset) / float(expand_w) # 计算边框x相对值
    gtbox[:, 1] = ((gtbox[:, 1] * image_h) + y_offset) / float(expand_h) # 计算边框y相对值
    gtbox[:, 2] = gtbox[:, 2] / x_ratio # 计算边框w相对值
    gtbox[:, 3] = gtbox[:, 3] / y_ratio # 计算边框h相对值
    
    return image, gtbox

def get_boxes_ious_xywh(box1, box2):
    """
    功能:
        计算边框列表的交并比
    输入:
        box1 - 边框列表1
        box2 - 边框列表2
    输出:
        ious - 交并比值列表
    """
    # 判断边框维度
    assert box1.shape[-1] == 4, "Box1 shape[-1] should be 4."
    assert box2.shape[-1] == 4, "Box2 shape[-1] should be 4."
    
    # 计算交集面积
    x1_min = box1[:, 0] - box1[:, 2]/2.0
    y1_min = box1[:, 1] - box1[:, 3]/2.0
    x1_max = box1[:, 0] + box1[:, 2]/2.0
    y1_max = box1[:, 1] + box1[:, 3]/2.0
    
    x2_min = box2[:, 0] - box2[:, 2]/2.0
    y2_min = box2[:, 1] - box2[:, 3]/2.0
    x2_max = box2[:, 0] + box2[:, 2]/2.0
    y2_max = box2[:, 1] + box2[:, 3]/2.0
    
    x_min = np.maximum(x1_min, x2_min)
    y_min = np.maximum(y1_min, y2_min)
    x_max = np.minimum(x1_max, x2_max)
    y_max = np.minimum(y1_max, y2_max)
    
    w = np.maximum(x_max - x_min, 0.0)
    h = np.maximum(y_max - y_min, 0.0)
    
    intersection = w * h # 交集面积
    
    # 计算并集面积
    s1 = box1[:, 2] * box1[:, 3]
    s2 = box2[:, 2] * box2[:, 3]
    
    union = s1 + s2 - intersection # 并集面积
    
    # 计算交并比值
    ious = intersection / union
    
    return ious

def get_cpbox_item(cpbox, gtbox, gtcls, image_size):
    """
    功能:
        计算裁剪边框
    输入:
        cpbox      - 裁剪边框: 真实位置
        gtbox      - 真实边框: 相对位置
        gtcls      - 真实类别
        image_size - 图像尺寸
    输出:
        gtbox      - 裁剪边框
        gtcls      - 裁剪类别
        box_number - 边框数量
    """
    # 拷贝真实边框列表
    gtbox = gtbox.copy() # 防止本次返回失败, 影响下次真实边框值
    gtcls = gtcls.copy() # 防止本次返回失败, 影响下次真实类别值
    
    # 转换真实边框位置:x1, y1, x2, y2
    image_w, image_h = map(float, image_size) # 读取图像宽高

    gtbox[:, 0], gtbox[:, 2] = \
        (gtbox[:, 0] - gtbox[:, 2]/2) * image_w, (gtbox[:, 0] + gtbox[:, 2]/2) * image_w # 先计算右边,再赋值左边,避免影响左边原始值
    gtbox[:, 1], gtbox[:, 3] = \
        (gtbox[:, 1] - gtbox[:, 3]/2) * image_h, (gtbox[:, 1] + gtbox[:, 3]/2) * image_h # 原址计算避免内存拷贝,分开计算,需要先拷贝gtbox
    
    # 转换裁剪边框位置
    cpbox_x, cpbox_y, cpbox_w, cpbox_h = map(float, cpbox) # 读取边框位置
    cpbox = np.array([cpbox_x, cpbox_y, cpbox_x + cpbox_w, cpbox_y + cpbox_h]) # 转换为ndarray格式:xyxy格式
    
    # 计算真实边框中心在裁剪边框内的个数的掩码
    gtbox_centers = (gtbox[:, :2] + gtbox[:, 2:]) / 2.0 # 计算所有真实边框的中心位置: x=(x1+x2)/2, y=(y1+y2)/2
    
    mask = np.logical_and(cpbox[:2] <= gtbox_centers, gtbox_centers <= cpbox[2:]).all(axis=1) # 中心位置掩码:真实边框中心是否在裁剪边框内
    
    # 计算裁剪边框位置
    gtbox[:, :2] = np.maximum(gtbox[:, :2], cpbox[:2]) # 计算真实边框与裁剪边框最大x1y1值
    gtbox[:, 2:] = np.minimum(gtbox[:, 2:], cpbox[2:]) # 计算真实边框与裁剪边框最小x2y2值
    gtbox[:, :2] -= cpbox[:2] # 计算真实边框x1y1值, 相对于裁剪边框的左上角位置
    gtbox[:, 2:] -= cpbox[:2] # 计算真实边框x2y2值, 相对于裁剪边框的左上角位置
    
    # 计算裁剪边框中左上角位置小于右下角的掩码
    mask = np.logical_and(mask, (gtbox[:, :2] < gtbox[:, 2:]).all(axis=1) ) # 裁剪边框的右下坐标是否在左上角左边之下
    
    # 计算裁剪边框位置
    gtbox = gtbox * np.expand_dims(mask.astype('float32'), axis=1) # 使用掩码计算合格的裁剪边框
    
    # 转换相对边框位置: x, y, w, h
    gtbox[:, 0], gtbox[:, 2] = \
        (gtbox[:, 0] + gtbox[:, 2])/2 / cpbox_w, (gtbox[:, 2] - gtbox[:, 0]) / cpbox_w # 先计算右边,再赋值左边,避免影响左边原始值
    gtbox[:, 1], gtbox[:, 3] = \
        (gtbox[:, 1] + gtbox[:, 3])/2 / cpbox_h, (gtbox[:, 3] - gtbox[:, 1]) / cpbox_h # 原址计算避免内存拷贝,分开计算,需要先拷贝gtbox
    
    # 计算裁剪边框类别
    gtcls = gtcls * mask.astype('float32') # 使用掩码计算合格的裁剪边框
    
    # 计算裁剪边框数量
    box_number = mask.sum() # 统计掩码为1的数量
    
    return gtbox, gtcls, box_number

def random_crop_image(image, gtbox, gtcls, crop_scope=[0.3, 1.0], max_ratio=2.0, iou_scopes=None, max_trial=50):
    """
    功能: 
        随机裁剪图像
    输入: 
        image      - 图像数据
        gtbox      - 边框列表
        gtcls      - 类别列表
        crop_scope - 裁剪比例范围
        max_ratio  - 最大裁剪比例
        iou_scopes - 交并比值范围
        max_trial  - 最大试验次数
    输出:
        image      - 图像数据
        gtbox      - 边框列表
        gtcls      - 类别列表
    """
    # 是否存在物体
    if random.random() > 0.5:
        return image, gtbox, gtcls
    if len(gtbox) == 0: # 如果物体边框数为零, 则返回原始图像数据项目
        return image, gtbox, gtcls
    
    # 交并比值范围
    if not iou_scopes: # 如果不存在交并比值范围列表, 则使用交默认并比值范围列表
        iou_scopes = [(0.1, 1.0), (0.3, 1.0), (0.5, 1.0), (0.7, 1.0), (0.9, 1.0), (0.0, 1.0)]
    
    # 转换图像格式
    image = Image.fromarray(image) # 转换为Image格式
    image_w, image_h = image.size # 获取图像宽高
    
    # 计算裁剪边框
    cpbox_list = [(0, 0, image_w, image_h)] # 裁剪边框列表
    
    for min_iou, max_iou in iou_scopes:
        for i in range(max_trial):
            # 随机生成裁剪比例
            crop_ratio = random.uniform(crop_scope[0], crop_scope[1]) # 随机生成裁剪宽高比例
            aspect_ratio = random.uniform(
                max(1 / max_ratio, (crop_ratio * crop_ratio)),
                min(max_ratio, 1 / (crop_ratio * crop_ratio))) # 随机生成裁剪纵横比例
            
            # 计算裁剪边框位置
            crop_h = int(image_h * crop_ratio / np.sqrt(aspect_ratio)) # 计算随机生成的裁剪高度
            crop_w = int(image_w * crop_ratio * np.sqrt(aspect_ratio)) # 计算随机生成的裁剪宽度
            
            crop_x = random.randint(0, image_w - crop_w) # 随机生成裁剪x坐标
            crop_y = random.randint(0, image_h - crop_h) # 随机生成裁剪y坐标
            
            # 计算边框交并集值
            cpbox = np.array(
                [[(crop_x + crop_w / 2.0) / float(image_w),
                  (crop_y + crop_h / 2.0) / float(image_h),
                  crop_w / float(image_w),
                  crop_h / float(image_h)]]) # 计算裁剪边框相对位置: xywh格式
            
            ious = get_boxes_ious_xywh(cpbox, gtbox) # 裁剪边框形状为(1, 4), 真实边框形状为(50, 4)
            
            # 添加裁剪边框列表
            if min_iou <= ious.min() and ious.max() <= max_iou: # 如果符合交并比值范围, 则添加一个裁剪边框, 并结束循环
                cpbox_list.append((crop_x, crop_y, crop_w, crop_h)) # 裁剪边框为真实位置
                break
                
    # 随机裁剪图像
    for i in range(len(cpbox_list)):
        # 弹出裁剪边框
        cpbox = cpbox_list.pop(random.randint(0, len(cpbox_list) - 1)) # 随机弹出裁剪边框
        
        # 计算裁剪边框
        crop_boxes, crop_gtcls, box_number = get_cpbox_item(cpbox, gtbox, gtcls, image.size)
        
        # 开始裁剪图像
        if box_number > 0: # 如果裁剪边框数量大于0, 则裁剪图像,并结束循环
            image = image.crop( (cpbox[0], cpbox[1], cpbox[0] + cpbox[2], cpbox[1] + cpbox[3]) ) # 用真实位置裁剪图像
            image = image.resize(image.size, Image.LANCZOS) # 高质量缩放到原来大小
            
            gtbox = crop_boxes # 返回裁剪边框
            gtcls = crop_gtcls # 返回裁剪类别
            
            break
    
    # 转换图像格式
    image = np.asarray(image) # 转换为ndarray格式
    
    return image, gtbox, gtcls

def random_interpolate_image(image, scale_size, interpolation=None):
    """
    功能: 
        随机插值图像
    输入: 
        image         - 图像数据
        scale_size    - 缩放宽高
        interpolation - 插值方法
    输出:
        image         - 图像数据
    """
    # 转换图像格式
    image = Image.fromarray(image) # 转换为Image格式
    
    # 随机缩放图像
    interpolation_method = [Image.NEAREST,Image.BILINEAR ,Image.BICUBIC, Image.LANCZOS] # 插值方法列表
    
    if not interpolation or interpolation not in interpolation_method:
        interpolation = interpolation_method[random.randint(0, len(interpolation_method) - 1)] # 随机选取插值方法
    
    image = image.resize(scale_size, interpolation)
    
    # 转换图像格式
    image = np.asarray(image) # 转换为ndarray格式，数据类型为HWC，uint8类型
    
    return image

def random_flip_image(image, gtbox):
    """
    功能: 
        随机翻转图像
    输入: 
        image - 图像数据
        gtbox - 边框列表
    输出:
        image - 图像数据
        gtbox - 边框列表
    """
    if random.random() > 0.5:
        image = image[:, ::-1, :] # 水平翻转图像, 列全部倒序排列
        gtbox[:, 0] = 1.0 - gtbox[:, 0] # 水平翻转边框, x坐标变为相反数
        
    return image, gtbox

def random_shuffle_gtbox(gtbox, gtcls):
    """
    功能: 
        随机打乱边框
    输入: 
        gtbox - 边框列表
        gtcls - 类别列表
    输出:
        gtbox - 边框列表
        gtcls - 类别列表
    """
    # 连接边框和类别
    data_list = np.concatenate([gtbox, gtcls[:, np.newaxis]], axis=1)
    
    # 打乱列表的顺序
    index = np.arange(data_list.shape[0])
    np.random.shuffle(index)
    
    data_list = data_list[index, :]
    
    # 保存边框和类别
    gtbox = data_list[:, :4]
    gtcls = data_list[:, -1]
    
    return gtbox, gtcls

def augment_image(image, gtbox, gtcls, scale_size):
    """
    功能:
        增强图像
    输入:
        image      - 图像数据
        gtbox      - 边框列表
        gtcls      - 类别列表
        scale_size - 缩放尺寸
    输出:
        image      - 图像数据
        gtbox      - 边框列表
        gtcls      - 类别列表
    """
    # 随机变换图像
    image = random_distort_image(image)
    
    # 随机填充图像
    image, gtbox = random_expand_image(image, gtbox)
    
    # 随机裁剪图像
    image, gtbox, gtcls = random_crop_image(image, gtbox, gtcls)
        
    # 随机插值图像
    image = random_interpolate_image(image, scale_size)
    
    # 随机翻转图像
    image, gtbox = random_flip_image(image, gtbox)
    
    # 随机打乱边框
    gtbox, gtcls = random_shuffle_gtbox(gtbox, gtcls)
    
    return image, gtbox, gtcls

##############################################################################################################

def get_scale_size(mode):
    """
    功能:
        获取缩放尺寸
    输入:
        mode       - 获取模式
    输出:
        scale_size - 缩放尺寸
    """
    if (mode == 'train') or (mode == 'valid'): # 如果是训练或验证模式, 随机生成缩放宽高
        scale_size = 320 + 32 * random.randint(0, 9) # 随机生成宽高,范围为[320, 608],步进为32
    else:
        scale_size = 608
    
    scale_size = (scale_size, scale_size) # 组合缩放宽高[w,h]
    
    return scale_size

def get_data_array(batch_data):
    """
    功能:
        将数据列表转换为数组构成的元组
    输入:
        batch_data       - 批次数据列表
    输出:
        image_array      - 图像数据数组
        gtbox_array      - 物体边框数组
        gtcls_array      - 图像类别数组
        image_size_array - 图像宽高数组
    """
    image_array = np.array([item[0] for item in batch_data], dtype='float32')
    gtbox_array = np.array([item[1] for item in batch_data], dtype='float32')
    gtcls_array = np.array([item[2] for item in batch_data], dtype='int32')
    image_size_array = np.array([item[3] for item in batch_data], dtype='int32')
    
    return image_array, gtbox_array, gtcls_array, image_size_array

def data_reader(data, scale_size):
    """
    功能:
        读取一条数据
    输入:
        data       - 数据项目
        scale_size - 缩放图像宽高
    输出:
        image      - 图像数据
        gtbox      - 边框列表
        gtcls      - 类别列表
        image_size - 原始图像高宽
    """
    # 读取数据
    image, gtbox, gtcls, image_size = get_data_item(data)

    # 增强图像
    image, gtbox, gtcls = augment_image(image, gtbox, gtcls, scale_size)
    
    # 减去均值
    mean = np.array([0.485, 0.456, 0.406]).reshape((1, 1, -1)) # COCO数据集通道平均值
    stdv = np.array([0.229, 0.224, 0.225]).reshape((1, 1, -1)) # COCO数据集通道标准差
    
    image = (image/255.0 - mean) / stdv # 对图像进行归一化
    image = image.astype('float32').transpose((2, 0, 1)) # 转换图片格式:[H,W,C]到[C,H,W]
    
    return image, gtbox, gtcls, image_size

def single_thread_reader(data_path, batch_size=8, mode='train'):
    """
    功能:
        单线程读取批次数据
    输入:
        data_path  - 数据集路径
        batch_size - 每批数据大小
        mode       - 读取数据模式: train或valid
    输出:
        reader     - 数据读取器
    """
    # 读取数据列表
    data_list = get_data_list(data_path)
    
    # 读取数据项目
    def reader():
        # 设置读取模式
        if mode == 'train': # 如果是训练模式, 则打乱数据列表
            np.random.shuffle(data_list)
        scale_size = get_scale_size(mode) # 设置缩放宽高
        
        # 输出批次数据
        batch_data = [] # 批次数据
        
        for item in data_list:
            # 读取数据
            image, gtbox, gtcls, image_size = data_reader(item, scale_size)
            
            # 输出数据
            batch_data.append((image, gtbox, gtcls, image_size))
            if len(batch_data) == batch_size: # 如果压入一批数据, 则弹出数据
                # 弹出数据
                yield get_data_array(batch_data)
                
                # 重置数据
                batch_data = []
                image_size = get_scale_size(mode)

        # 输出剩余数据
        if len(batch_data) > 0:
            yield get_data_array(batch_data)
    
    return reader

def multip_thread_reader(data_path, batch_size=8, mode='train'):
    """
    功能:
        多线程读取批次数据
    输入:
        data_path  - 数据集路径
        batch_size - 每批数据大小
        mode       - 读取数据模式: train 或 valid
    输出:
        reader     - 数据读取器
    """
    # 读取数据列表
    data_list = get_data_list(data_path)
    
    # 读取数据项目
    def item_loader():
        # 设置读取模式
        if mode == 'train': # 如果是训练模式, 则打乱数据列表
            np.random.shuffle(data_list)
        scale_size = get_scale_size(mode) # 设置缩放宽高
        
        # 输出批次项目
        batch_item = [] # 批次项目
        
        for item in data_list:
            # 输出项目
            batch_item.append((item, scale_size))
            if len(batch_item) == batch_size: # 如果压入一批项目, 则弹出项目
                # 弹出数据
                yield batch_item
                
                # 重置项目
                batch_item = []
                image_size = get_scale_size(mode)

        # 输出剩余数据
        if len(batch_item) > 0:
            yield batch_item
    
    # 读取数据内容
    def data_loader(batch_item):
        batch_data = [] # 批次数据
        
        for item, scale_size in batch_item:
            image, gtbox, label, image_size = data_reader(item, scale_size)
            batch_data.append((image, gtbox, label, image_size))
        
        return get_data_array(batch_data)
    
    # 多线程读取器
    reader = paddle.reader.xmap_readers(data_loader, item_loader, process_num=4, buffer_size=16)
    
    return reader

##############################################################################################################

def single_test_reader(image_path, scale_size=(608, 608)):
    """
    功能:
        读取一张预测图像
    输入:
        image_path - 图像路径
        scale_size - 缩放高宽
    输出:
        image      - 图像数据
        image_size - 原始高宽
    """
    # 读取图像
    image = Image.open(image_path)                   # 读取图像
    if image.mode != 'RGB':
        image = image.convert('RGB')
    image_size = (image.size[0], image.size[1])      # 读取尺寸
    image = image.resize(scale_size, Image.BILINEAR) # 缩放图像
    
    # 转换格式
    image = np.array(image, dtype='float32')
    image_size = np.array(image_size, dtype='int32')
    
    # 减去均值
    mean = np.array([0.485, 0.456, 0.406]).reshape((1, 1, -1)) # COCO数据集通道平均值
    stdv = np.array([0.229, 0.224, 0.225]).reshape((1, 1, -1)) # COCO数据集通道标准差
    
    image = (image/255.0 - mean) / stdv                  # 对图像进行归一化
    image = image.astype('float32').transpose((2, 0, 1)) # 转换图片格式:[H,W,C]到[C,H,W]
    
    # 增加维度
    image = np.expand_dims(image, axis=0) # 增加维度: (1,3,608,608)
    image_size = np.expand_dims(image_size, axis=0) # 增加维度: (1,2)
    
    return image, image_size

def display_infer(infer, image_path):
    """
    功能:
        显示预测结果
    输入:
        infer       - 预测结果
        image_path  - 图像路径
    输出:
    """
    # 读取图像
    image = Image.open(image_path)   # 读取图像
    if image.mode != 'RGB':
        image = image.convert('RGB') # 转换格式
    
    # 绘制结果
    object_names = ['Boerner','Leconte','Linnaeus','acuminatus','armandi','coleoptera','linnaeus'] # 物体名称
    color = ['r', 'g', 'b', 'c','m', 'y', 'k']                                                     # 边框颜色
    
    plt.figure(figsize=(10, 10)) # 设置显示图像大小
    currentAxis = plt.gca()      # 获取图像当前坐标
    
    for item in infer: # 遍历预测结果
        # 获取结果
        index = int(item[0])        # 类别索引
        names = object_names[index] # 类别名称
        pdbox = item[2:6]           # 边框位置
        
        # 绘制边框
        rectangle = patches.Rectangle(                                       # 设置边框
            (pdbox[0], pdbox[1]), pdbox[2]-pdbox[0]+1, pdbox[3]-pdbox[1]+1, 
            linewidth=1, edgecolor=color[index], facecolor=color[index], 
            fill=False, linestyle='-')
        currentAxis.add_patch(rectangle)                                     # 绘制边框
        plt.text(pdbox[0], pdbox[1], names, fontsize=12, color=color[index]) # 绘制类别
    
    # 显示结果
    plt.imshow(image)
    plt.show()

def get_test_array(batch_data):
    """
    功能:
        将数据列表转换为数组构成的元组
    输入:
        batch_data       - 数据列表
    输出:
        image_name_array - 图像名字数组
        image_array      - 图像数据数组
        image_size_array - 图像宽高数组
    """
    image_name_array = np.array([item[0] for item in batch_data])
    image_array = np.array([item[1] for item in batch_data], dtype='float32')
    image_size_array = np.array([item[2] for item in batch_data], dtype='int32')
    
    return image_name_array, image_array, image_size_array

def multip_test_reader(data_path, batch_size=1, scale_size=(608, 608)):
    """
    功能:
        读取一批测试数据
    输入:
        data_path  - 数据目录
        batch_size - 每批大小
        scale_size - 缩放高宽
    输出:
        image_name - 图像名称
        image      - 图像数据
        image_size - 原始高宽
    """
    # 读取数据列表
    data_list = os.listdir(data_path)
    
    # 读取一批数据
    def reader():
        # 输出批次数据
        batch_data = [] # 批次数据
        
        for image_name in data_list:
            # 读取图像路径
            image_path = os.path.join(data_path, image_name)
            
            # 读取一张图像
            image = Image.open(image_path)                   # 读取图像
            if image.mode != 'RGB':
                image = image.convert('RGB')
            image_size = (image.size[0], image.size[1])      # 图像大小
            image = image.resize(scale_size, Image.BILINEAR) # 缩放图像
            
            # 转换格式
            image = np.array(image, dtype='float32')
            image_size = np.array(image_size, dtype='int32')

            # 减去均值
            mean = np.array([0.485, 0.456, 0.406]).reshape((1, 1, -1)) # COCO数据集通道平均值
            stdv = np.array([0.229, 0.224, 0.225]).reshape((1, 1, -1)) # COCO数据集通道标准差

            image = (image/255.0 - mean) / stdv                  # 对图像进行归一化
            image = image.astype('float32').transpose((2, 0, 1)) # 转换图片格式:[H,W,C]到[C,H,W]
            
            # 输出数据
            batch_data.append((image_name.split('.')[0], image, image_size))
            if len(batch_data) == batch_size: # 如果压入一批数据, 则弹出数据
                # 弹出数据
                yield get_test_array(batch_data)
                
                # 重置数据
                batch_data = []

        # 输出剩余数据
        if len(batch_data) > 0:
            yield get_test_array(batch_data)
    
    return reader